from pyasn1_modules.rfc2251 import SearchResultEntry, BindRequest, AuthenticationChoice, BindResponse, SearchRequest, LDAPMessage, PartialAttributeList, SearchResultDone, AttributeValue, Attribute, AttributeValue, ModifyResponse
from pyasn1.type.univ import Sequence, SetOf
from pyasn1.type import univ
from pyasn1.type.tag import Tag
from pyasn1.error import SubstrateUnderrunError
from pyasn1.codec.ber import decoder, encoder
import sqlite3
import cPickle as pickle
from utils import config, GServer, GAnswer
import copy

def get_type(o):
	l = o.prettyPrint()
	po = l.find('protocolOp=')
	po = l[po + 11 : l.find(':', po)]
	if po != 'Choice':
		return po
	fc = l.find('Choice:') + 10
	po = l[fc : l.find('=', fc)]
	return po

# LDAP Server with sqlite3 base
class LdapServer(GServer):
	def getdb(self, dn):
		if not self.dn_case_sensitive:
			dn = dn.upper()
		c = self.db.cursor()
		c.execute('select data from data where id=:id limit 1;', {'id': dn})
		d = {}
		for row in c:
			d = pickle.loads(str(row[0]))
		c.close()
		return d

	def setdb(self, dn, d):
		if not self.dn_case_sensitive:
			dn = dn.upper()
		c = self.db.cursor()
		s = pickle.dumps(d)
		c.execute('UPDATE data SET data=:data WHERE id=:id;', {'data': s, 'id': dn})
		c.close()
		self.db.commit()

	def populate(self, cfg_name):
		self.log('populate %s' % cfg_name)
		self.db = sqlite3.connect(":memory:")
		cur = self.db.cursor()
		cur.execute("create table data(id text primary key, data blob)")
		c = config.get(cfg_name)
		c = c['data']
		load = 'populate:\n'
		for e in c:
			if 'seq' in e and 'template' in e:
				seq = config.get_obj(e['seq'])
				t = e['template']
				for x in seq.gen():
					v = copy.deepcopy(t)
					v = v.popitem()
					dc = v[0].replace(seq.var, x)
					if not self.dn_case_sensitive:
						dc = dc.upper()
					for j in v[1]:
						v[1][j] = str(v[1][j]).replace(seq.var, str(x))
					s = pickle.dumps(v[1])
					cur.execute('INSERT INTO data VALUES (?,?);', (dc, s))
				load += '%s : %s\n' % (repr(seq), t)
		cur.close()
		self.db.commit()
		return GAnswer(True, load[:-1])

	def handler(self, socket, address):
		ldap = LDAPMessage()
		while True:
			ber = ''
			b = None
			size = 65536
			while True:
				try:
					ber += socket.recv(size)
				except Exception, e:
					self.log_error(e)
					socket.close()
					return None
				if ber == '':
					socket.close()
					return None
			
				ldap = LDAPMessage()
				try:
					b, rest = decoder.decode(ber, asn1Spec = ldap)
				except SubstrateUnderrunError, e:
					e = str(e)
					e = e[:e.find('-')]
					size = int(e)
				else:
					break
			self.log_debug('got: %s' % b.prettyPrint())
			msid = b.getComponentByName('messageID')
			t = get_type(b)
			if t == 'bindRequest':
				self.log('BindRequest id:%s' % msid)
				r = BindResponse()
				r.setComponentByPosition(0, 0)
				r.setComponentByPosition(1, '')
				r.setComponentByPosition(2, '')
				b.setComponentByPosition(1, r)
				socket.sendto(encoder.encode(b), address)
			elif t == 'searchRequest':
				dn = str(b.getComponentByName('protocolOp').getComponentByName('searchRequest').getComponentByName('baseObject'))
				attr = b.getComponentByName('protocolOp').getComponentByName('searchRequest').getComponentByName('attributes')
				attr_len = len(attr._componentValues)
				self.log('SearchRequest id:%s attr_cnt:%i dn:%s ' % (msid, attr_len, dn))
				d = self.getdb(dn)
				r = SearchResultEntry()
				r.setComponentByPosition(0, dn)
				cnt = 0
				e = PartialAttributeList()
				rc = 0 # success
				b = LDAPMessage()
				b.setComponentByName('messageID', msid)
				attrs = []
				nattrs = []
				if len(d) == 0:
					rc = 32 # noSuchObject
					self.log_error('dn: "%s" not found' % dn)
				else:
					if attr_len == 0:
						for i in d:
							av = AttributeValue(d[i])
							so = SetOf()
							so.setComponentByPosition(0, av)
							a = Attribute()
							a.setComponentByPosition(0, i)
							a.setComponentByPosition(1, so)
							e.setComponentByPosition(cnt, a)
							cnt += 1
					else:
						for i in range(attr_len):
							v = str(attr.getComponentByPosition(i))
							attrs.append(v)
							if v in d:
								av = AttributeValue(d[v])
								so = SetOf()
								so.setComponentByPosition(0, av)
								a = Attribute()
								a.setComponentByPosition(0, v)
								a.setComponentByPosition(1, so)
								e.setComponentByPosition(cnt, a)
								cnt += 1
							else:
								#pass
								nattrs.append(v)		
					r.setComponentByPosition(1, e)
					b.setComponentByPosition(1, r)
					a = self.send_chunking(encoder.encode(b), socket)
					self.loga(a)
				if len(attrs) != 0:
					self.log('required attrs: %s' % attrs)
				if len(nattrs) != 0:
					self.log('missed attrs: %s' % nattrs)
				b = LDAPMessage()
				r = SearchResultDone()
				r.setComponentByPosition(0, rc)
				r.setComponentByPosition(1, '')
				r.setComponentByPosition(2, '')
				b.setComponentByPosition(1, r)
				b.setComponentByName('messageID', msid)
				socket.sendto(encoder.encode(b), address)
			elif t =='modifyRequest':
				#dn = str(b.getComponentByName('protocolOp').getComponentByName('modifyRequest').getComponentByName('baseObject'))
				dn = str(b.getComponentByPosition(1).getComponentByPosition(7).getComponentByPosition(0))
				self.log('ModifyRequest id:%s dn:%s' % (msid, dn))
				m = b.getComponentByPosition(1).getComponentByPosition(7).getComponentByPosition(1)
				attr_len = len(m._componentValues)
				d = self.getdb(dn)
				s = '\n'
				for i in range(attr_len):
					n = m.getComponentByPosition(i)
					n = n.getComponentByName('modification')
					v = str(n.getComponentByPosition(1).getComponentByPosition(0))
					d[str(n.getComponentByPosition(0))] = v
					if len(v) > 32:
						s += '%s = %s ...\n' % (str(n.getComponentByPosition(0)), repr(v[:32]))
					else:
						s += '%s = %s\n' % (str(n.getComponentByPosition(0)), v)
				self.log('modification list:%s' % s[:-1])
				self.setdb(dn, d)
				r = ModifyResponse()
				r.setComponentByPosition(0, self.modify_resp_code)
				r.setComponentByPosition(1, '')
				r.setComponentByPosition(2, '')
				b.setComponentByPosition(1, r)
				socket.sendto(encoder.encode(b), address)
			elif t == 'unbindRequest':
				self.log('UnbindRequest id:%s' % msid)
				break
			else:
				self.log_error('unhandled %s' % t)

